package com.whut.wesocket;

import com.whut.handler.SshSessionHandler;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.websocket.OnClose;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Created by YY on 2018-02-02.
 */
@Component
@ServerEndpoint(value="/ssh/{key}")
public class SshWebSocketServlet {

    //key: username:password@address
    private static final Map<String,Set<Session>> clientMap = new ConcurrentHashMap<>();
    private static final Map<String,SshSessionHandler> handlerMap = new ConcurrentHashMap<>();

    @OnOpen
    public void onOpen(@PathParam("key") String key, Session session) {
        if (clientMap.get(key)==null){
            Set<Session> sessions = new HashSet<Session>();
            sessions.add(session);
            clientMap.put(key,sessions);
        }else {
            clientMap.get(key).add(session);
        }
        SshSessionHandler sessionHandler =new SshSessionHandler();
        sessionHandler.setWebSocketServlet(this,key);
        handlerMap.put(key,sessionHandler);
        Thread thread =new Thread(sessionHandler);
        thread.start();
    }

    @OnMessage
    public void onMessage(String message,@PathParam("key") String key) {
        handlerMap.get(key).setUserInput(message);
    }

    @OnClose
    public void onClose(@PathParam("key") String key, Session peer) {
        Set<Session> sessions = clientMap.get(key);
        synchronized (sessions){
            sessions.remove(peer);
        }
        if (sessions.size()<=0) {
            clientMap.remove(key);
        }
        handlerMap.get(key).stopHandler();
        handlerMap.remove(key);
    }

    public boolean sendMes(String key,String message) {
        try {
            Set<Session> sessions = clientMap.get(key);
            if (sessions==null) {
                return true;
            }
            synchronized (sessions){
                for (Session session : sessions) {
                    session.getBasicRemote().sendText(message);
                }
            }
            return true;
        }catch (Exception e){
            e.printStackTrace();
            return false;
        }
    }
}
